import rclpy
from rclpy.node import Node
import cv2
import numpy as np
import serial
import time
import math
from sensor_msgs.msg import Image
from cv_bridge import CvBridge

# -------------------------- 协议常量定义（与STM32匹配） --------------------------
FRAME_HEADER = 0x66  # 帧头
FRAME_LENGTH = 0x07  # 帧长度（固定7字节）
FRAME_TAIL = 0x5B    # 帧尾
SEND_INTERVAL = 0.05  # 发送间隔（20Hz，避免数据拥堵）

# 以下函数保持不变
def init_camera_params():
    """初始化双目相机参数（包含畸变矫正参数）"""
    left_camera_matrix = np.array([[715.088662856028, -1.27750597113329, 321.725929157012],
                                   [0, 717.938092436680, 251.352481762661],
                                   [0., 0., 1.]])
    right_camera_matrix = np.array([[712.537452801351, -1.49459241948369, 325.232106968860],
                                    [0, 715.552769066380, 250.629910089625],
                                    [0., 0., 1.]])

    left_distortion = np.array(
        [[-0.440986151974959, 0.327464329751247, -0.000547150225973162, -0.00126940700042557, -0.335074894313998]])
    right_distortion = np.array(
        [[-0.444122749137210, 0.348225581098782, 0.00176792559350367, 0.00123611244757822, -0.660993977721134]])

    R = np.array([[0.999866011965132, 0.000416637643297468, 0.0163641232584165],
                  [-0.000544744807817963, 0.999969236927380, 0.00782486114692518],
                  [-0.0163603597159969, -0.00783272698033651, 0.999835479975588]])
    T = np.array([-59.8387197927031, 0.0107452720401373, 0.330883446596301])
    size = (640, 480)

    # 计算立体校正参数
    R1, R2, P1, P2, Q, validPixROI1, validPixROI2 = cv2.stereoRectify(
        left_camera_matrix, left_distortion,
        right_camera_matrix, right_distortion,
        size, R, T
    )

    # 生成左右相机的矫正映射表（用于后续图像矫正）
    left_map1, left_map2 = cv2.initUndistortRectifyMap(
        left_camera_matrix, left_distortion, R1, P1, size, cv2.CV_16SC2)
    right_map1, right_map2 = cv2.initUndistortRectifyMap(
        right_camera_matrix, right_distortion, R2, P2, size, cv2.CV_16SC2)

    return {
        "left_camera_matrix": left_camera_matrix,
        "right_camera_matrix": right_camera_matrix,
        "left_distortion": left_distortion,
        "right_distortion": right_distortion,
        "R": R,
        "T": T,
        "size": size,
        "Q": Q,
        "left_map1": left_map1,
        "left_map2": left_map2,
        "right_map1": right_map1,
        "right_map2": right_map2,
        "baseline": abs(T[0])
    }

def init_stereo_matcher(blockSize=5, numDisparities=128):
    """初始化SGBM匹配器和WLS滤波器（用于视差计算）"""
    img_channels = 3
    stereo = cv2.StereoSGBM_create(
        minDisparity=0,
        numDisparities=numDisparities,
        blockSize=blockSize,
        P1=8 * img_channels * blockSize * blockSize,
        P2=32 * img_channels * blockSize * blockSize,
        disp12MaxDiff=5,
        preFilterCap=63,
        uniquenessRatio=15,
        speckleWindowSize=200,
        speckleRange=32,
        mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY
    )

    right_matcher = cv2.ximgproc.createRightMatcher(stereo)
    wls_filter = cv2.ximgproc.createDisparityWLSFilter(matcher_left=stereo)
    wls_filter.setLambda(8000)
    wls_filter.setSigmaColor(1.5)

    return stereo, right_matcher, wls_filter

def init_serial(port='COM7', baudrate=115200):
    """初始化串口通信（与STM32通信）"""
    try:
        ser = serial.Serial(
            port=port,
            baudrate=baudrate,
            bytesize=serial.EIGHTBITS,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            timeout=0.1
        )
        print(f"串口已打开: {ser.port}")
        return ser
    except serial.SerialException as e:
        print(f"串口打开失败: {e}")
        print("将以无串口模式运行")
        return None

def rectify_image(img, map1, map2):
    """对图像进行畸变矫正（支持彩色图和灰度图）"""
    return cv2.remap(img, map1, map2, cv2.INTER_LINEAR)

def compute_disparity_and_depth(left_rect, right_rect, stereo, right_matcher, wls_filter, Q, depth_scale=16.0):
    """计算视差图和深度信息"""
    disparity_left = stereo.compute(left_rect, right_rect)
    disparity_right = right_matcher.compute(right_rect, left_rect)

    filtered_disparity = wls_filter.filter(disparity_left, left_rect, None, disparity_right)
    filtered_disparity = np.clip(filtered_disparity, 0, None)

    # 视差图归一化（用于显示）
    disp_norm = cv2.normalize(filtered_disparity, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
    disp_color = cv2.applyColorMap(disp_norm, cv2.COLORMAP_JET)

    # 计算三维坐标（深度信息）
    threeD = cv2.reprojectImageTo3D(filtered_disparity, Q, handleMissingValues=False)
    threeD = threeD * depth_scale  # 缩放深度值

    return filtered_disparity, disp_norm, disp_color, threeD

def original_to_rectified(x, y, left_map1):
    """将原始图像坐标映射到矫正后坐标（用于对齐检测结果）"""
    h, w = left_map1.shape[:2]
    if not (0 <= x < w and 0 <= y < h):
        return None
    # 从映射表中获取矫正后坐标（CV_16SC2格式：每个元素是[x_rect, y_rect]）
    map_val = left_map1[y, x]
    x_rect = int(map_val[0])
    y_rect = int(map_val[1])
    if 0 <= x_rect < w and 0 <= y_rect < h:
        return (x_rect, y_rect)
    return None

def detect_orange_ball(rect_color, threeD, left_map1):
    """在矫正后的彩色图上检测橙色球（确保坐标与深度图对齐）"""
    # 转换到HSV颜色空间进行颜色检测
    hsv = cv2.cvtColor(rect_color, cv2.COLOR_BGR2HSV)
    lower_orange = np.array([15, 82, 210])  # 橙色HSV范围（适配不同光照）
    upper_orange = np.array([68, 255, 255])
    mask = cv2.inRange(hsv, lower_orange, upper_orange)

    # 形态学操作去除噪声
    kernel = np.ones((5, 5), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)  # 闭合孔洞
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)   # 去除小噪点

    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    ball_info = None
    best_score = 0

    for contour in contours:
        area = cv2.contourArea(contour)
        if area < 50:  # 过滤小轮廓
            continue

        # 计算圆形度（筛选圆形物体）
        perimeter = cv2.arcLength(contour, True)
        if perimeter == 0:
            continue
        circularity = 4 * math.pi * area / (perimeter **2)
        if circularity < 0.6:  # 只保留接近圆形的物体
            continue

        # 计算轮廓中心（矫正后坐标）
        (cx_rect, cy_rect), radius = cv2.minEnclosingCircle(contour)
        cx_rect, cy_rect = int(cx_rect), int(cy_rect)

        # 验证深度有效性
        if 0 <= cx_rect < threeD.shape[1] and 0 <= cy_rect < threeD.shape[0]:
            depth = threeD[cy_rect, cx_rect]
            z_val = depth[2]
            if z_val <= 0:
                continue

            # 计算实际距离（米）
            distance = np.linalg.norm(depth) / 1000.0

            # 评分（优先选择面积大、圆形度高的目标）
            score = area * circularity
            if score > best_score:
                best_score = score
                # 将矫正后坐标映射回原始坐标（用于发送给下位机）
                orig_coord = None  # 如需原始坐标可在此处反向映射
                ball_info = {
                    "rect_center": (cx_rect, cy_rect),  # 矫正后坐标（用于绘图）
                    "orig_center": orig_coord,          # 原始坐标（如需可开启）
                    "depth": depth,
                    "distance": distance,
                    "contour": contour
                }

    return ball_info

def send_coordinate(ser, x, y):
    """按协议发送坐标给下位机"""
    if not ser or not ser.is_open:
        return

    # 限制坐标范围为0~65535（16位无符号整数）
    x_clamped = np.clip(int(x), 0, 65535)
    y_clamped = np.clip(int(y), 0, 65535)

    # 拆分高低字节
    x_low = x_clamped & 0xFF
    x_high = (x_clamped >> 8) & 0xFF
    y_low = y_clamped & 0xFF
    y_high = (y_clamped >> 8) & 0xFF

    # 构建帧结构：[帧头, 长度, X低, X高, Y低, Y高, 帧尾]
    frame = bytearray([
        FRAME_HEADER, FRAME_LENGTH,
        x_low, x_high, y_low, y_high,
        FRAME_TAIL
    ])

    # 发送数据
    try:
        ser.write(frame)
        print(f"发送坐标: ({x_clamped}, {y_clamped})  帧: {[hex(b) for b in frame]}")
    except serial.SerialException as e:
        print(f"串口发送失败: {e}")

def onmouse_pick_points(event, x, y, flags, param):
    """鼠标回调：点击深度图获取三维坐标（用于调试）"""
    threeD = param
    if threeD is None or event != cv2.EVENT_LBUTTONDOWN:
        return

    if 0 <= y < threeD.shape[0] and 0 <= x < threeD.shape[1]:
        depth = threeD[y, x]
        z_val = depth[2]
        if z_val <= 0:
            print("无效深度值")
            return
        xyz = depth / 1000.0  # 转换为米
        distance = np.linalg.norm(xyz)
        print(f"点击坐标: ({x}, {y})  三维坐标: ({xyz[0]:.3f}, {xyz[1]:.3f}, {xyz[2]:.3f})m  距离: {distance:.3f}m")


class StereoCameraNode(Node):
    def __init__(self):
        super().__init__('stereo_camera_node')
        # 初始化相机参数和组件
        self.cam_params = init_camera_params()
        self.stereo, self.right_matcher, self.wls_filter = init_stereo_matcher()
        self.ser = init_serial()
        self.screen_center = (self.cam_params["size"][0] // 2, self.cam_params["size"][1] // 2)  # 图像中心点
        self.get_logger().info(f"图像中心点: {self.screen_center}")

        # 打开相机（双目相机通常输出1280x480分辨率，左右图各640x480）
        self.cap = cv2.VideoCapture(0)
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        if not self.cap.isOpened():
            self.get_logger().error("无法打开相机")
            return

        # 创建显示窗口
        cv2.namedWindow("Rectified Color", cv2.WINDOW_NORMAL)  # 矫正后的彩色图（主显示）
        cv2.namedWindow("Original Color", cv2.WINDOW_NORMAL)   # 原始彩色图
        cv2.namedWindow("Original Gray", cv2.WINDOW_NORMAL)    # 原始灰度图
        cv2.namedWindow("Depth Map", cv2.WINDOW_NORMAL)        # 深度图（用于调试）
        cv2.setMouseCallback("Depth Map", onmouse_pick_points)

        self.last_send_time = time.time()
        self.bridge = CvBridge()

        # 创建图像发布者
        self.rectified_color_pub = self.create_publisher(Image, 'rectified_color_image', 10)
        self.original_color_pub = self.create_publisher(Image, 'original_color_image', 10)
        self.original_gray_pub = self.create_publisher(Image, 'original_gray_image', 10)
        self.depth_map_pub = self.create_publisher(Image, 'depth_map_image', 10)

        # 创建定时器
        self.timer = self.create_timer(0.1, self.process_frame)

    def process_frame(self):
        ret, frame = self.cap.read()
        if not ret:
            self.get_logger().error("无法读取帧，退出")
            return

        # 分割左右相机图像
        frame_left = frame[:, :640]  # 左相机原始彩色图
        frame_right = frame[:, 640:] # 右相机原始彩色图

        # 生成原始灰度图
        gray_left = cv2.cvtColor(frame_left, cv2.COLOR_BGR2GRAY)

        # 对左相机图像进行畸变矫正（彩色图和灰度图）
        left_rect_gray = rectify_image(gray_left, self.cam_params["left_map1"], self.cam_params["left_map2"])
        left_rect_color = rectify_image(frame_left, self.cam_params["left_map1"], self.cam_params["left_map2"])  # 矫正后的彩色图
        right_rect_gray = rectify_image(cv2.cvtColor(frame_right, cv2.COLOR_BGR2GRAY),
                                       self.cam_params["right_map1"], self.cam_params["right_map2"])

        # 计算视差和深度
        _, _, disp_color, threeD = compute_disparity_and_depth(
            left_rect_gray, right_rect_gray, self.stereo, self.right_matcher, self.wls_filter, self.cam_params["Q"]
        )

        # 在矫正后的彩色图上检测橙色球
        ball_info = detect_orange_ball(left_rect_color, threeD, self.cam_params["left_map1"])

        # 绘制检测结果（在矫正后的彩色图上）
        target_x, target_y = self.screen_center  # 默认发送中心点坐标
        if ball_info:
            # 获取检测信息
            cx_rect, cy_rect = ball_info["rect_center"]
            distance = ball_info["distance"]
            contour = ball_info["contour"]
            target_x, target_y = cx_rect, cy_rect  # 更新发送坐标为小球中心

            # 绘制小球轮廓
            cv2.drawContours(left_rect_color, [contour], -1, (0, 255, 0), 2)
            # 绘制小球中心
            cv2.circle(left_rect_color, (cx_rect, cy_rect), 5, (0, 0, 255), -1)
            # 绘制与中心点的连线
            cv2.line(left_rect_color, self.screen_center, (cx_rect, cy_rect), (0, 255, 255), 2)
            # 显示坐标和距离（基于彩色图圆心坐标）
            cv2.putText(left_rect_color, f"Center: ({cx_rect}, {cy_rect})", (10, 30),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
            cv2.putText(left_rect_color, f"Distance: {distance:.3f}m", (10, 60),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
        else:
            # 未检测到目标时显示提示
            cv2.putText(left_rect_color, "No Target", (10, 30),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

        # 绘制图像中心点标记
        cv2.circle(left_rect_color, self.screen_center, 5, (255, 0, 0), -1)
        cv2.line(left_rect_color, (self.screen_center[0]-10, self.screen_center[1]),
                (self.screen_center[0]+10, self.screen_center[1]), (255, 0, 0), 2)
        cv2.line(left_rect_color, (self.screen_center[0], self.screen_center[1]-10),
                (self.screen_center[0], self.screen_center[1]+10), (255, 0, 0), 2)

        # 按间隔发送坐标（20Hz）
        current_time = time.time()
        if current_time - self.last_send_time > SEND_INTERVAL:
            send_coordinate(self.ser, target_x, target_y)
            self.last_send_time = current_time

        # 显示所有要求的图像
        cv2.imshow("Rectified Color", left_rect_color)  # 矫正后的彩色图（带标注）
        cv2.imshow("Original Color", frame_left)        # 原始彩色图
        cv2.imshow("Original Gray", gray_left)          # 原始灰度图
        cv2.imshow("Depth Map", disp_color)             # 深度图（调试用）

        # 发布图像消息
        rectified_color_msg = self.bridge.cv2_to_imgmsg(left_rect_color, encoding="bgr8")
        original_color_msg = self.bridge.cv2_to_imgmsg(frame_left, encoding="bgr8")
        original_gray_msg = self.bridge.cv2_to_imgmsg(gray_left, encoding="mono8")
        depth_map_msg = self.bridge.cv2_to_imgmsg(disp_color, encoding="bgr8")

        self.rectified_color_pub.publish(rectified_color_msg)
        self.original_color_pub.publish(original_color_msg)
        self.original_gray_pub.publish(original_gray_msg)
        self.depth_map_pub.publish(depth_map_msg)

        # 退出条件（按q或ESC）
        key = cv2.waitKey(1)
        if key in [ord('q'), 27]:
            self.cap.release()
            cv2.destroyAllWindows()
            if self.ser and self.ser.is_open:
                self.ser.close()
                self.get_logger().info("串口已关闭")
            self.get_logger().info("程序退出")
            rclpy.shutdown()


def main(args=None):
    rclpy.init(args=args)
    node = StereoCameraNode()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()


if __name__ == "__main__":
    main()
